package com.such.kit.web.jwt;

import java.util.Map;
import org.jose4j.jwa.AlgorithmConstraints;
import org.jose4j.jwa.AlgorithmConstraints.ConstraintType;
import org.jose4j.jwk.HttpsJwks;
import org.jose4j.jwk.JsonWebKeySet;
import org.jose4j.jwk.RsaJsonWebKey;
import org.jose4j.jwk.RsaJwkGenerator;
import org.jose4j.jws.AlgorithmIdentifiers;
import org.jose4j.jws.JsonWebSignature;
import org.jose4j.jwt.JwtClaims;
import org.jose4j.jwt.consumer.InvalidJwtException;
import org.jose4j.jwt.consumer.JwtConsumer;
import org.jose4j.jwt.consumer.JwtConsumerBuilder;
import org.jose4j.keys.resolvers.HttpsJwksVerificationKeyResolver;
import org.jose4j.keys.resolvers.JwksVerificationKeyResolver;
import org.jose4j.lang.JoseException;
import com.such.kit.Callable;
import com.such.kit.Logger;
import com.such.kit.cache.CacheCommons;
import com.such.kit.validate.ValidateCommons;
import com.such.kit.web.jwt.bean.JWTConfig;
import com.such.kit.web.jwt.bean.ValidateResult;

/**
 * @author SUCH
 * <pre>
 * JWT 通用工具类
 * </pre>
 */
public class JWTCommons {

	private static String jwtCache = CacheCommons.buildCacheName(JWTCommons.class);

	private JWTConfig jwtConfig;
	private RsaJsonWebKey rsaJWK;

	/**
	 * @param jwtConfig JWT 配置
	 * @throws JoseException
	 */
	private JWTCommons(JWTConfig jwtConfig) throws JoseException {
		this.jwtConfig = jwtConfig;
		this.rsaJWK = RsaJwkGenerator.generateJwk(2048);
		this.rsaJWK.setKeyId(jwtConfig.getKey());
	}

	/**
	 * <pre>
	 * 创建 token
	 * </pre>
	 * @param payload 自定义载荷
	 * @return token
	 * @throws JoseException
	 */
	public String generate(Map<String, Object> payload) throws JoseException {
		String subject = this.jwtConfig.getSubject();
		String issuer = this.jwtConfig.getIssuer();
		String audience = this.jwtConfig.getAudience();
		int expAtNow = this.jwtConfig.getExpAtNow();
		int nbfAtNow = this.jwtConfig.getNbfAtNow();

		JwtClaims claims = new JwtClaims();
		claims.setGeneratedJwtId(); // jti 唯一身份标识
		claims.setIssuedAtToNow(); // iat 签发时间
		if (ValidateCommons.isNotEmpty(subject)) {
			claims.setSubject(subject); // sub 面向的用户（主题）
		}
		if (ValidateCommons.isNotEmpty(issuer)) {
			claims.setIssuer(issuer); // iss 签发者
		}
		if (ValidateCommons.isNotEmpty(audience)) {
			claims.setAudience(audience); // aud 接收者
		}
		if (expAtNow != 0) {
			claims.setExpirationTimeMinutesInTheFuture(expAtNow); // exp 有效期
		}
		if (nbfAtNow != 0) {
			claims.setNotBeforeMinutesInThePast(nbfAtNow); // nbf 不可用时间（指定时间之后可用）
		}
		if (ValidateCommons.isNotEmpty(payload)) {
			for (String key : payload.keySet()) {
				Object value = payload.get(key);
				claims.setClaim(key, value);
			}
		}

		JsonWebSignature jws = new JsonWebSignature();
		jws.setKey(this.rsaJWK.getPrivateKey());
		jws.setKeyIdHeaderValue(this.rsaJWK.getKeyId());
		jws.setAlgorithmHeaderValue(AlgorithmIdentifiers.RSA_USING_SHA256);
		jws.setPayload(claims.toJson());
		return jws.getCompactSerialization();
	}

	/**
	 * <pre>
	 * 验证 token
	 * </pre>
	 * @param token 令牌
	 * @return 验证结果
	 */
	public ValidateResult validate(String token) {
		String subject = this.jwtConfig.getSubject();
		String issuer = this.jwtConfig.getIssuer();
		String audience = this.jwtConfig.getAudience();
		boolean useJWKS = this.jwtConfig.useJWKS();
		String jwksEndpoint = this.jwtConfig.getJwksEndpoint();
		
		JwtConsumerBuilder builder = new JwtConsumerBuilder();
		builder.setJwsAlgorithmConstraints(new AlgorithmConstraints(ConstraintType.WHITELIST, AlgorithmIdentifiers.RSA_USING_SHA256));
		builder.setRequireExpirationTime(); // 启用令牌期限检查
		builder.setAllowedClockSkewInSeconds(3000);
		if (ValidateCommons.isNotEmpty(subject)) {
			builder.setExpectedSubject(subject);
		}
		if (ValidateCommons.isNotEmpty(issuer)) {
			builder.setExpectedIssuer(issuer);
		}
		if (ValidateCommons.isNotEmpty(audience)) {
			builder.setExpectedAudience(audience);
		}
		if (useJWKS) {
			if (ValidateCommons.isNotEmpty(jwksEndpoint)) {
				HttpsJwks httpsJkws = new HttpsJwks(jwksEndpoint);
				HttpsJwksVerificationKeyResolver httpsJwksKeyResolver = new HttpsJwksVerificationKeyResolver(httpsJkws);
				builder.setVerificationKeyResolver(httpsJwksKeyResolver); // 使用 HTTPS JWKS 站点公钥
			} else {
				JsonWebKeySet jsonWebKeySet = new JsonWebKeySet(this.rsaJWK);
				JwksVerificationKeyResolver jwksResolver = new JwksVerificationKeyResolver(jsonWebKeySet.getJsonWebKeys());
				builder.setVerificationKeyResolver(jwksResolver); // 使用系统生成的公钥
			}
		} else {
			builder.setVerificationKey(this.rsaJWK.getKey()); // 签名
		}

		ValidateResult result = null;
		try {
			JwtConsumer jwtConsumer = builder.build();
			JwtClaims jwtClaims = jwtConsumer.processToClaims(token);
			result = new ValidateResult(jwtClaims);
		} catch (InvalidJwtException e) {
			result = new ValidateResult(e);
		}
		return result;
	}

	/**
	 * @param jwtConfig JWT 配置
	 * @return JWTCommons
	 */
	public static JWTCommons getInstance(final JWTConfig jwtConfig) {
		JWTCommons jwtCommons = null;
		try {
			jwtCommons = CacheCommons.get(jwtCache, jwtConfig.getKey(), new Callable<String, JWTCommons>() {
				@Override
				public JWTCommons call(String cacheName) throws Exception {
					return new JWTCommons(jwtConfig);
				}
			});
		} catch (Exception e) {
			Logger.error(JWTCommons.class, "创建 JWTCommons 异常", e);
		}
		return jwtCommons;
	}

}
